import torch
import torch.nn as nn

from monai.metrics import DiceMetric, SurfaceDiceMetric, HausdorffDistanceMetric


class Metrics(nn.Module):
    def __init__(self, num_classes) -> None:
        super().__init__()
        
        self.dice = DiceMetric(include_background=True,
                               num_classes=num_classes)

        self.nsd = SurfaceDiceMetric(include_background=True,
                               class_thresholds=[4, 4, 2])

        self.hd = HausdorffDistanceMetric(include_background=True,
                               percentile=95)
        
    def forward(self, x, y):
        dice = self.dice(x, y)
        nsd = self.nsd(x, y)
        hd = torch.nan_to_num(self.hd(x, y), 0.)
        
        return dice, nsd, hd
        

class MyDice(nn.Module):
    def __init__(self, n_class):
        super().__init__()
        self.n_class = n_class

    def compute_stats_brats(self, p, y):
        scores = torch.zeros(self.n_class, device=p.device, dtype=torch.float32)
        p = (torch.sigmoid(p) > 0.5).int()

        for i in range(self.n_class):
            p_i, y_i = p[:, i], y[:, i]
            if (y_i != 1).all():
                # no foreground class
                scores[i - 1] += 1 if (p_i != 1).all() else 0
                continue
            tp, fn, fp = self.get_stats(p_i, y_i, 1)
            denom = (2 * tp + fp + fn).to(torch.float)
            score_cls = (2 * tp).to(torch.float) / denom if torch.is_nonzero(denom) else 0.0
            scores[i] += score_cls
        return scores

    @staticmethod
    def get_stats(p, y, c):
        tp = torch.logical_and(p == c, y == c).sum()
        fn = torch.logical_and(p != c, y == c).sum()
        fp = torch.logical_and(p == c, y != c).sum()
        return tp, fn, fp

    def forward(self, p, y):
        dice += self.compute_stats_brats(p, y) if self.brats else self.compute_stats(p, y)
        return dice
    

if __name__ == "__main__":
    import numpy as np
    def trans_brats_label(x):
            mask_WT = x.copy()
            mask_WT[mask_WT == 1] = 1
            mask_WT[mask_WT == 2] = 1
            mask_WT[mask_WT == 3] = 1

            mask_TC = x.copy()
            mask_TC[mask_TC == 1] = 1
            mask_TC[mask_TC == 2] = 0
            mask_TC[mask_TC == 3] = 1

            mask_ET = x.copy()
            mask_ET[mask_ET == 1] = 0
            mask_ET[mask_ET == 2] = 0
            mask_ET[mask_ET == 3] = 1
            
            mask = np.stack([mask_WT, mask_TC, mask_ET], axis=0)
            return mask
    
    path = '/Users/qlc/Desktop/Dataset/Brats2023/Adult_Glioma/TrainingData/BraTS-GLI-00127-000/BraTS-GLI-00127-000-seg.nii.gz'
    
    import SimpleITK as sitk
    label = sitk.ReadImage(path)
    label = sitk.GetArrayFromImage(label)
    label = trans_brats_label(label)

    import torch
    label = torch.tensor(label)


    pred = label[:, 35].unsqueeze(0)
    true = label[:, 45].unsqueeze(0)
    
    monai_dice = Metrics(num_classes=3)
    
    out, a, b = monai_dice(pred, true)
    print(out, a, b)
    b = torch.nan_to_num(b, 0.)
    print(b)

    
